// csp.go
package main

import (
	"errors"
	"log"
	"strconv"

	pbc "github.com/Nik-U/pbc"
)

// CSP 代表云服务提供商
type CSP struct {
	ID        string
	Tags      []*pbc.Element // 存储认证标签
	FileID    string
	params    *SystemParams
	dataOwner *DataOwner // 模拟拥有对DO原始数据的访问权限（仅用于测试）
}

// NewCSP 创建一个新的CSP实例
func NewCSP(id string, params *SystemParams, do *DataOwner) *CSP {
	return &CSP{ID: id, params: params, dataOwner: do}
}

// StoreTags 存储数据所有者上传的认证标签
func (c *CSP) StoreTags(fileID string, tags []*pbc.Element) {
	c.FileID = fileID
	c.Tags = tags
}

// PartialProofGen 核心算法：生成部分证明 Γ_k = (α_k, β_k)
func (c *CSP) PartialProofGen(challenge *AuditChallenge) (*PartialProof, error) {
	if challenge.FileID != c.FileID {
		return nil, errors.New("file ID mismatch")
	}

	alpha_k := c.params.Pairing.NewZr().Set0()
	beta_k := c.params.Pairing.NewG1().Set1()
	originalBlocks := c.dataOwner.originalData[c.FileID]

	for i, idx := range challenge.Indices {
		if idx >= len(originalBlocks) || idx >= len(c.Tags) {
			return nil, errors.New("index out of bounds")
		}
		weight := challenge.Weights[i]
		block_data := originalBlocks[idx]
		tag := c.Tags[idx]

		// α_k = Σ e_l * x_kil
		alpha_k.Add(alpha_k, c.params.Pairing.NewZr().Mul(weight, block_data))
		// β_k = Π (y_kil)^e_l
		beta_k.Mul(beta_k, c.params.Pairing.NewG1().PowZn(tag, weight))
	}

	log.Printf("[%s] Generated partial proof (α_k, β_k).", c.ID)
	return &PartialProof{Alpha: alpha_k, Beta: beta_k}, nil
}

// AggregateProof 聚合所有部分证明 (模拟)
func (c *CSP) AggregateProof(partialProofs []*PartialProof) *FinalProof {
	finalAlpha := c.params.Pairing.NewZr().Set0()
	finalBeta := c.params.Pairing.NewG1().Set1()

	for _, p := range partialProofs {
		finalAlpha.Add(finalAlpha, p.Alpha)
		finalBeta.Mul(finalBeta, p.Beta)
	}
	log.Printf("[%s] Aggregated all partial proofs to get final proof (α, β).", c.ID)
	return &FinalProof{Alpha: finalAlpha, Beta: finalBeta}
}

// VerifyProof 核心算法：验证最终证明 (已修正)
func (c *CSP) VerifyProof(finalProof *FinalProof, challenge *AuditChallenge, pk *PublicKey, cspIDs []string) bool { // 接收 cspIDs
	log.Printf("[%s] Starting verification...", c.ID)
	// LHS (不变)
	lhs := c.params.Pairing.NewGT().Pair(finalProof.Beta, c.params.G)

	// RHS
	h_agg := c.params.Pairing.NewG1().Set1()
	for _, cspID := range cspIDs { // 使用传入的 cspIDs
		for i, idx := range challenge.Indices {
			weight := challenge.Weights[i]
			h_val := c.params.h(cspID + "||" + challenge.FileID + "||" + strconv.Itoa(idx))
			h_agg.Mul(h_agg, c.params.Pairing.NewG1().PowZn(h_val, weight))
		}
	}

	u_alpha := c.params.Pairing.NewG1().PowZn(c.params.U, finalProof.Alpha)
	rhs_base := c.params.Pairing.NewG1().Mul(h_agg, u_alpha)
	rhs := c.params.Pairing.NewGT().Pair(rhs_base, pk.G_gamma)

	result := lhs.Equals(rhs)
	log.Printf("[%s] Verification result: %t", c.ID, result)
	return result
}
